# 使用vllm进行单机多卡推理
import re
from datasets import load_dataset
from vllm import LLM, SamplingParams
import multiprocessing as mp

# ========== 模型路径（必须是 HF 格式或 trust_remote_code 支持）==========
model_path = "/home/ma-user/work/emapo/Qwen3-4B1/checkpoint-1650"


# ========== 提取数值函数（用于准确率评估）==========
def extract_number(text, is_pred=False):
    def answer_split(t):
        return t.split("nswer:")[-1] if "nswer:" in t else t

    target = answer_split(text) if is_pred else text
    match = re.search(r"[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?", target)
    return float(match.group()) if match else None


# ========== 评估准确率函数 ==========
def compute_accuracy(completions, prompts, references):
    cnt = 0
    for i in range(len(prompts)):
        print(f"\nidx: {i} --------------\nprompt: {prompts[i]}\completions: {completions[i]}\nreference: {references[i]}")
        if references[i] in completions[i]:
            cnt += 1
    return cnt / len(prompts)


# ========== 数据集预处理 ==========
def preprocess_dataset(path, name):
    if name == "GSM8K":
        ds = load_dataset(path, "main", cache_dir=f"/home/ma-user/work/DownLoads/Dataset/{path}")
    else:
        ds = load_dataset(path, cache_dir=f"/home/ma-user/work/DownLoads/Dataset/{path}")
    if name in ["GSM8K", "MinervaMath", "AMC-2023", "MATH-500"]:
        dataset = ds["test"]
    elif name in ["AIME-2024", "AIME-2025"]:
        dataset = ds["train"]
    else:
        raise ValueError(f"Unsupported dataset name: {name}")

    def get_processor(name):
        if name == "GSM8K":
            return lambda e: {"prompt": e["question"], "reference": e["answer"].split("####")[-1].strip()}
        elif name == "MinervaMath":
            return lambda e: {"prompt": e["question"], "reference": e["answer"]}
        elif name == "AMC-2023":
            return lambda e: {"prompt": e["question"], "reference": str(e["answer"])}
        elif name == "MATH-500":
            return lambda e: {"prompt": e["problem"], "reference": str(e["answer"])}
        elif name == "AIME-2024":
            return lambda e: {"prompt": e["Problem"], "reference": str(e["Answer"])}
        elif name == "AIME-2025":
            return lambda e: {"prompt": e["problem"], "reference": str(e["answer"])}
        else:
            raise ValueError(f"No processor for {name}")

    processor = get_processor(name)
    dataset = dataset.map(processor, remove_columns=dataset.column_names, num_proc=64)
    return dataset


# ========== 使用 vLLM 执行推理 ==========
def test_model_vllm(
    llm,
    dataset,
    name,
    max_new_tokens=5000,
    temperature=0.6,
    top_p=0.95,
    top_k=20,
):
    print(f"\n===== [vLLM] Evaluating on {name} =====")

    prompts = dataset["prompt"]
    references = dataset["reference"]

    sampling_params = SamplingParams(
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        max_tokens=max_new_tokens,
        stop=["</s>"],  # 可自定义终止符号
    )

    # 批量推理
    outputs = llm.generate(prompts, sampling_params)

    completions = [output.outputs[0].text for output in outputs]

    acc = compute_accuracy(completions, prompts, references)
    print(f"[{name}] Accuracy: {acc:.3f} ({int(acc * len(prompts))} / {len(prompts)})")
    return acc


# ========== 主程序入口 ==========
if __name__ == "__main__":
    mp.set_start_method("spawn", force=True)
    test_sets = {
        "GSM8K": "openai/gsm8k",
        "MinervaMath": "math-ai/minervamath",
        "AMC-2023": "zwhe99/amc23",
        "MATH-500": "HuggingFaceH4/MATH-500",
        "AIME-2024": "Maxwell-Jia/AIME_2024",
        "AIME-2025": "MathArena/aime_2025",
    }
    # 初始化 vLLM 模型（自动选择 GPU，支持多卡）
    llm = LLM(model=model_path, trust_remote_code=True, dtype="auto", tensor_parallel_size=4)
    results = {}
    for name, path in test_sets.items():
        try:
            processed_ds = preprocess_dataset(path, name)
            acc = test_model_vllm(llm, processed_ds, name)
            results[name] = acc
        except Exception as e:
            print(f"[{name}] Failed to evaluate: {e}")

    print("\n===== Summary =====")
    for name, acc in results.items():
        print(f"{name}: {acc:.3f}")
